#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Kalman filter utilities migrated from legacy implementation.
"""

import logging
from typing import Optional, Tuple

import numpy as np

logger = logging.getLogger("RDK_YOLO")


class KalmanFilter:
    """卡尔曼滤波器，用于目标运动预测。"""

    def __init__(self, dt: float = 0.1, process_noise_scale: float = 0.01, measurement_noise_scale: float = 0.1):
        self.state = np.zeros((4, 1))
        self.initialized = False

        self.F = np.array([[1, 0, dt, 0], [0, 1, 0, dt], [0, 0, 1, 0], [0, 0, 0, 1]])
        self.H = np.array([[1, 0, 0, 0], [0, 1, 0, 0]])
        self.Q = np.array(
            [
                [dt**4 / 4, 0, dt**3 / 2, 0],
                [0, dt**4 / 4, 0, dt**3 / 2],
                [dt**3 / 2, 0, dt**2, 0],
                [0, dt**3 / 2, 0, dt**2],
            ]
        ) * process_noise_scale
        self.R = np.eye(2) * measurement_noise_scale
        self.P = np.eye(4)

        self.x_pred = None
        self.P_pred = None
        self.K = None

        self.current_velocity = np.zeros(2)
        self.last_position: Optional[Tuple[float, float]] = None
        self.last_time: Optional[float] = None

        logger.debug(
            "卡尔曼滤波器初始化: dt=%.3f, 过程噪声=%.3f, 测量噪声=%.3f",
            dt,
            process_noise_scale,
            measurement_noise_scale,
        )

    def init(self, position: Tuple[float, float], time_stamp: float) -> None:
        self.state = np.array([[position[0]], [position[1]], [0], [0]])
        self.last_position = position
        self.last_time = time_stamp
        self.initialized = True
        logger.debug("卡尔曼滤波器初始化状态: 位置=(%.1f, %.1f)", position[0], position[1])

    def update(self, position: Tuple[float, float], time_stamp: float):
        if not self.initialized:
            self.init(position, time_stamp)
            return position, (0.0, 0.0)

        dt = time_stamp - self.last_time if self.last_time is not None else 0.0
        if dt <= 0:
            dt = 0.001

        self.F[0, 2] = dt
        self.F[1, 3] = dt

        self.Q[0, 0] = dt**4 / 4
        self.Q[0, 2] = dt**3 / 2
        self.Q[1, 1] = dt**4 / 4
        self.Q[1, 3] = dt**3 / 2
        self.Q[2, 0] = dt**3 / 2
        self.Q[2, 2] = dt**2
        self.Q[3, 1] = dt**3 / 2
        self.Q[3, 3] = dt**2

        self.x_pred = self.F @ self.state
        self.P_pred = self.F @ self.P @ self.F.T + self.Q

        S = self.H @ self.P_pred @ self.H.T + self.R
        self.K = self.P_pred @ self.H.T @ np.linalg.inv(S)

        measurement = np.array([[position[0]], [position[1]]])
        y = measurement - self.H @ self.x_pred
        self.state = self.x_pred + self.K @ y
        self.P = (np.eye(4) - self.K @ self.H) @ self.P_pred

        self.current_velocity = (self.state[2, 0], self.state[3, 0])

        self.last_position = position
        self.last_time = time_stamp

        estimated_position = (self.state[0, 0], self.state[1, 0])
        estimated_velocity = self.current_velocity

        logger.debug(
            "卡尔曼滤波器更新: 测量=(%.1f, %.1f), 估计=(%.1f, %.1f), 速度=(%.1f, %.1f), dt=%.3f",
            position[0],
            position[1],
            estimated_position[0],
            estimated_position[1],
            estimated_velocity[0],
            estimated_velocity[1],
            dt,
        )

        return estimated_position, estimated_velocity

    def predict(self, dt: Optional[float] = None):
        if not self.initialized:
            return None

        if dt is None:
            dt = self.F[0, 2]

        F_pred = np.array([[1, 0, dt, 0], [0, 1, 0, dt], [0, 0, 1, 0], [0, 0, 0, 1]])
        future_state = F_pred @ self.state

        predicted_position = (future_state[0, 0], future_state[1, 0])

        logger.debug("卡尔曼滤波器预测: dt=%.3f, 预测位置=(%.1f, %.1f)", dt, predicted_position[0], predicted_position[1])
        return predicted_position

    def reset(self) -> None:
        self.state = np.zeros((4, 1))
        self.P = np.eye(4)
        self.initialized = False
        self.last_position = None
        self.last_time = None
        self.current_velocity = np.zeros(2)
        logger.debug("卡尔曼滤波器重置")
